/*
 * elem_nmf_main.cpp
 *
 * Copyright 2013 Long Tran (ltran3@gatech.edu)
 *
 * main file for maxtrix factorization program
 * Solving 0.5|A-WH^T|^2 + <Lambda, W-U> + <Pi, H-V> + rho/2 (|W-U|^2 + |H-V|^2)
 *  Generate data
 *   mpirun -np 16 ./nmf  --gen --data=m1000_10_200_16 --m=200 --n=200 --k=10 
 *  Sequential NMF (ALS)
 *   mpirun -np 1 ./nmf  --anls --m=200 --n=200 --k=10 --iter=1000 --rho=10
 *  Master/Slave NMF
 *   mpirun -np 4 ./nmf --anls-dist --k=10 --data=m1000_10_200_16 --rho=14 --iter=5000 --interval=1000
 *  Decentralized NMF (specify network)
 *   mpirun -np 4 ./nmf --anls-dist --k=10 --data=m1000_10_200_16 --rho=14 --graph=grid --iter=5000 --interval=1000
 *   mpirun -np 4 ./nmf --anls-dist --k=10 --data=m1000_10_200_16 --rho=14 --graph=equal --degree=2 --iter=5000 --interval=1000
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#include <armadillo>
#include "helper/helper.h"
#include <elemental.hpp>
#include "algorithm/StopWatch.h"

using namespace elem;
typedef double R;
typedef elem::Complex<R> C;

po::variables_map vm;
po::options_description desc("Allowed options");

void process_program_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce help message")
    ("verbose", "print all debug information")
    ("interval", po::value<int>()->default_value(100), "data dimension hint")
    // DATA
    ("data", po::value<std::string>(), "data file")
    ("output", po::value<std::string>(), "output basis file")
    ("iterinfo", po::value<std::string>(), "iteration information file")
    // MODEL
    ("model", po::value<std::string>(), "output model file")
    // METHODS
    ("gen", "generate random data (MPI_SIZE)\n--m --n --k")
    ("anls-seq", "solve A = WH^T by Sequential Alternating Nonnegative Least Square\n--data --rate --rho --iter --k")
    ("anls-dist", "solve A = WH^T by Distributed Alternating Nonnegative Least Square\n--data --rate --rho --graph [--degree] --iter --k")
    ("lf-dist", "solve A = UV^T by Distributed Low-Rank Matrix Factorization\n--data --alpha --rho --graph [--degree] --iter --k")
    ("image", "read images")
    // COMMON PARAMETERS
    ("iter", po::value<int>()->default_value(1000), "maximum iterations")
    ("rate", po::value<double>()->default_value(1.0), "learning rate / step size")
    ("density", po::value<double>()->default_value(0.3), "sparsity density")
    ("rho", po::value<double>()->default_value(0.1), "augmented penalty")
    ("alpha", po::value<double>()->default_value(0.1), "concensus penalty")
    ("normalize", "keep columns of H normalized (norm <= 1)")
    ("l1", po::value<double>(), "l1 regularization")
    ("l2", po::value<double>(), "l2 regularization")
    // DISTRIBUTED ALGORITHMS
    ("graph", po::value<std::string>(), "decentralized ADMM with bipartie graph (grid|equal)")
    ("degree", po::value<int>(), "decentralized ADMM with bipartie graph")
    // GENERATE DATA
    ("m", po::value<int>()->default_value(10), "= A.nrows")
    ("n", po::value<int>()->default_value(10), "= A.ncols")
    ("k", po::value<int>()->default_value(2), "= W.ncols = H.nrows")
    ("np", po::value<int>()->default_value(2), "Number of processes to generate data")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm); 
}

class ElemANLS {
public:
  mpi::Comm& m_Comm;
  Grid& m_Grid;
  int m_CommRank;

  DistMatrix<R> *pA, W, H, E;
  int m, n, k;
  double rho;

  DistMatrix<R> I;
  DistMatrix<R> U, V, Lambda, Pi;
  DistMatrix<R> Temp,  Temp1, Temp2;
  DistMatrix<R> Temp3, Temp4, Temp5;

  // iteration information
  struct Iteration
  {
    int iter;
    double residual;
    TimeDuration duration, compute, communication;
    Iteration(int iter_, double residual_) : iter(iter_), residual(residual_) {}
  };
  std::vector<Iteration> iterInfo;
public:
  ElemANLS(mpi::Comm& comm, Grid& grid) 
    : m_Comm(comm), m_Grid(grid), m_CommRank(mpi::CommRank(comm)),
      pA(NULL), W(m_Grid), H(m_Grid), E(m_Grid), I(m_Grid),
      U(m_Grid), V(m_Grid), Lambda(m_Grid), Pi(m_Grid),
      Temp(m_Grid), Temp1(m_Grid), Temp2(m_Grid),
      Temp3(m_Grid), Temp4(m_Grid), Temp5(m_Grid)
  {
    PO_VARIABLE_VALUE(rho, vm, "rho", double);
  }

  void SetVariables(DistMatrix<R>& A, int k)
  {
    this->pA = &A;
    this->m = A.Height();
    this->n = A.Width();
    this->k = k;

    
    Uniform(m,k,U); Uniform(m,k,Lambda); 
    Uniform(n,k,V); Uniform(n,k,Pi); 
    Uniform(k,k,I);
    Zero(Lambda); 
    Zero(Pi); 
    Zero(I);
    for (int i = 0; i < k; i++) I.Set(i,i,rho);

    Uniform(m,k,Temp); Uniform(k,m,Temp1); Uniform(k,k,Temp2);
    Uniform(n,k,Temp3); Uniform(k,n,Temp4); Uniform(k,k,Temp5);
  }
  
  void UpdateW()
  {
    DistMatrix<R>& A = *pA;
    // W = (AH - Lambda + rho U) ( H^T H + rho I_k)^{-1}
    Zero(Temp);
    Axpy(R(rho), U, Temp); // rho U
    Axpy(R(-1), Lambda, Temp); // -Lambda
    Gemm(NORMAL, NORMAL, R(1), A, H, R(1), Temp); // + AH
    Transpose(Temp, Temp1);

    Temp2 = I;
    Gemm(TRANSPOSE, NORMAL, R(1), H, H, R(1), Temp2);
    CholeskySolve(UPPER, Temp2, Temp1);
    Transpose(Temp1, W);
  }

  void UpdateH()
  {
    DistMatrix<R>& A = *pA;
    // H = (A^T W - Pi + rho V) (W^T W + rho I_k)^{-1}
    Zero(Temp3);
    Axpy(R(rho), V, Temp3); // rho V
    Axpy(R(-1), Pi, Temp3); // -Pi
    Gemm(TRANSPOSE, NORMAL, R(1), A, W, R(1), Temp3); // +A^T W
    Transpose(Temp3, Temp4);
    
    Temp5 = I;
    Gemm(TRANSPOSE, NORMAL, R(1), W, W, R(1), Temp5);
    CholeskySolve(UPPER, Temp5, Temp4);
    Transpose(Temp4, H);
  }

  void UpdateUV()
  {
    // U = P_+(W + Lambda/rho)
    U = W;
    Axpy(R(1.0/rho), Lambda, U);
    for (int i = 0; i < U.LocalHeight(); i++) 
      for (int j = 0; j < U.LocalWidth(); j++)
        if (U.Get(i,j) < 1e-16) U.SetLocal(i,j,R(0));

    // V = P_+(H + Pi/rho)
    V = H;
    Axpy(R(1.0/rho), Pi, V);
    for (int i = 0; i < V.LocalHeight(); i++) 
      for (int j = 0; j < V.LocalWidth(); j++)
        if (V.Get(i,j) < 1e-16) V.SetLocal(i,j,R(0));
  }

  void UpdateLambdaPi()
  {
    // Lambda
    Temp=W;
    Axpy(R(-1), U, Temp);    // W-U
    Axpy(R(rho), Temp, Lambda); // Lambda += rho*(W-U)

    // Pi
    Temp3=H;
    Axpy(R(-1), V, Temp3);    // H-V
    Axpy(R(rho), Temp3, Pi); // Pi += rho*(H-V)
  }

  void SolveANLS()
  {
    StopWatch watch;
    TimeDuration runtime;
    
    DistMatrix<R>& A = *pA;
    Uniform(m, k, W);
    Uniform(n, k, H);
    PO_VARIABLE(maxiter, vm, "iter", int);
    PO_VARIABLE(interval, vm, "interval", int);

    iterInfo.clear();
    for (int iter = 0; iter < maxiter; iter++) {
      watch.GetDuration();
      UpdateW();
      UpdateH();
      UpdateUV();
      UpdateLambdaPi();
      runtime += watch.GetDuration();

      if ( iter == 0 || iter == maxiter-1 || iter % interval == 0 ) {
	E = A;
	Gemm(NORMAL, TRANSPOSE, R(-1), U, V, R(1), E);
	const R frobNormOfE = Norm( E, FROBENIUS_NORM );
	const R frobNormOfU = Norm( U, FROBENIUS_NORM );
	const R frobNormOfV = Norm( V, FROBENIUS_NORM );
	if ( m_CommRank == 0 ) {
	  Iteration info(iter, frobNormOfE);
	  info.duration = runtime;
	  iterInfo.push_back(info);

	  std::cerr << "iter=" << iter 
		    << " |E|=" << frobNormOfE
		    << " |U|=" << frobNormOfU
		    << " |V|=" << frobNormOfV
		    << " d=" << (double) info.duration.total_milliseconds()/1000
		    << std::endl;

	}
      }
    }//for
  }

};

int main(int argc, char** argv)
{
  Initialize( argc, argv );
  mpi::Comm comm = mpi::COMM_WORLD;
  const int commRank = mpi::CommRank( comm );

  process_program_options(argc, argv);

  try {
    PO_VARIABLE(m, vm, "m", int);
    PO_VARIABLE(n, vm, "n", int);
    PO_VARIABLE(k, vm, "k", int);
    PO_VARIABLE(maxiter, vm, "iter", int);
    PO_VARIABLE(interval, vm, "interval", int);
    PO_VARIABLE(rho, vm, "rho", double);

    PO_VARIABLE(data, vm, "data", std::string);
    std::stringstream ss;
    ss << data << "_" << 0; // commRank;
    arma::mat Aread;
    Aread.load(ss.str(), arma::arma_binary);
    m = Aread.n_rows; n = Aread.n_cols;
    
    const int nb = 96;
    const bool print = false; 
    //ProcessInput();
    //PrintInputReport();

    SetBlocksize( nb );
    Grid g( comm );

    DistMatrix<R> A( m, n, g ), W0(g), H0(g), E(m,n,g);

    /*
    // generate data
    Uniform( m, k, W0 );
    Uniform( n, k, H0 );
    Zero(A);
    Gemm( NORMAL, TRANSPOSE, R(1), W0, H0, R(1), A);
    */
    for (int i=0; i < m; i++)
      for (int j=0; j < n; j++)
	A.Set(i,j, Aread(i,j));

    /*
    E = A;
    Gemm(NORMAL, TRANSPOSE, R(-1), W0, H0, R(1), E);
    const R frobNormOfE = Norm( E, FROBENIUS_NORM );
    const R frobNormOfU = Norm( W0, FROBENIUS_NORM );
    const R frobNormOfV = Norm( H0, FROBENIUS_NORM );
    const R frobNormOfA = Norm( A, FROBENIUS_NORM );
    if ( commRank == 0 ) {
      std::cerr << "ElemANLS Ground truth |A|=" << frobNormOfA
		<< " |E|=" << frobNormOfE
		<< " |W|=" << frobNormOfU
		<< " |H|=" << frobNormOfV
		<< std::endl;
    }
    */

    if( print )
      A.Print("A");
    ElemANLS a(comm,g);
    a.SetVariables(A, k);
    a.SolveANLS();
  }
  catch ( std::exception& e ) {
    std::cerr << "Rank=" << commRank << " error: " << e.what() << std::endl;
  }
  Finalize();
  return 0;
}
